﻿/* System.Compute: a Computation Framework for .NET/Mono
 * Copyright (C) 2011 Ralf "raboni" Boeckhorst
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 3 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General
 * Public License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin Street, Fifth Floor,
 * Boston, MA 02110-1301 USA
 */

using System.Runtime.InteropServices;
using System.Text;

namespace System.Compute.OpenCL
{
    public sealed class Kernel : ICloneable, IDisposable
    {
        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern IntPtr clCreateKernel(IntPtr program, string kernel_name, ref ErrorCode errcode_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clCreateKernelsInProgram(IntPtr program, uint num_kernels, IntPtr[] kernels, ref uint num_kernels_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clCreateKernelsInProgram(IntPtr program, uint num_kernels, IntPtr[] kernels, uint num_kernels_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clRetainKernel(IntPtr kernel);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clReleaseKernel(IntPtr kernel);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clSetKernelArg(IntPtr kernel, uint arg_index, SizeT arg_size, IntPtr arg_value);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelInfo(IntPtr kernel, KernelInfo param_name, SizeT param_value_size, IntPtr param_value, ref SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelInfo(IntPtr kernel, KernelInfo param_name, SizeT param_value_size, byte[] param_value, SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelInfo(IntPtr kernel, KernelInfo param_name, SizeT param_value_size, ref uint param_value, SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelInfo(IntPtr kernel, KernelInfo param_name, SizeT param_value_size, ref IntPtr param_value, SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelWorkGroupInfo(IntPtr kernel, IntPtr device, KernelWorkGroupInfo param_name, SizeT param_value_size, ref SizeT param_value, SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelWorkGroupInfo(IntPtr kernel, IntPtr device, KernelWorkGroupInfo param_name, SizeT param_value_size, SizeT[] param_value, SizeT param_value_size_ret);

        [DllImport("opencl.dll", ExactSpelling = true)]
        static extern ErrorCode clGetKernelWorkGroupInfo(IntPtr kernel, IntPtr device, KernelWorkGroupInfo param_name, SizeT param_value_size, ref ulong param_value, SizeT param_value_size_ret);

        public string FunctionName
        {
            get
            {
                SizeT param_value_size = 0;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.FunctionName, 0, IntPtr.Zero, ref param_value_size);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                byte[] param_value = new byte[(long)param_value_size];
                res = clGetKernelInfo(Handle, KernelInfo.FunctionName, param_value_size, param_value, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return Encoding.Default.GetString(param_value);
            }
        }

        public uint NumArgs
        {
            get
            {
                uint ret = 0;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.NumArgs, sizeof(uint), ref ret, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return ret;
            }
        }

        public uint ReferenceCount
        {
            get
            {
                uint ret = 0;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.ReferenceCount, sizeof(uint), ref ret, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return ret;
            }
        }

        public Context Context
        {
            get
            {
                IntPtr tmp = IntPtr.Zero;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.Context, IntPtr.Size, ref tmp, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return new Context(tmp);
            }
        }

        public Program Program
        {
            get
            {
                IntPtr tmp = IntPtr.Zero;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.Program, IntPtr.Size, ref tmp, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return new Program(tmp);
            }
        }

        public string Attributes
        {
            get
            {
                SizeT param_value_size = 0;
                ErrorCode res = clGetKernelInfo(Handle, KernelInfo.Attributes, 0, IntPtr.Zero, ref param_value_size);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                byte[] param_value = new byte[(long)param_value_size];
                res = clGetKernelInfo(Handle, KernelInfo.Attributes, param_value_size, param_value, 0);
                if (res != ErrorCode.Success)
                    throw new OpenCLException(res);
                return Encoding.Default.GetString(param_value);
            }
        }

        internal IntPtr Handle;

        internal Kernel(IntPtr handle)
        {
            Handle = handle;
        }

        public Kernel(Program program, string kernel_name)
        {
            ErrorCode res = ErrorCode.Success;
            Handle = clCreateKernel(program.Handle, kernel_name, ref res);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
        }

        public void SetArgument(uint arg_index, DeviceBuffer buffer)
        {
            GCHandle h = GCHandle.Alloc(buffer.Handle, GCHandleType.Pinned);
            ErrorCode res = clSetKernelArg(Handle, arg_index, IntPtr.Size, h.AddrOfPinnedObject());
            h.Free();
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
        }

        public void SetArgument(uint arg_index, DeviceImage image)
        {
            GCHandle h = GCHandle.Alloc(image.Handle, GCHandleType.Pinned);
            ErrorCode res = clSetKernelArg(Handle, arg_index, IntPtr.Size, h.AddrOfPinnedObject());
            h.Free();
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
        }

        public void SetArgument(uint arg_index, DeviceSampler sampler)
        {
            GCHandle h = GCHandle.Alloc(sampler.Handle, GCHandleType.Pinned);
            ErrorCode res = clSetKernelArg(Handle, arg_index, IntPtr.Size, h.AddrOfPinnedObject());
            h.Free();
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
        }

        public void SetArgument<T>(uint arg_index, T value) where T : struct
        {
            GCHandle h = GCHandle.Alloc(value, GCHandleType.Pinned);
            ErrorCode res = clSetKernelArg(Handle, arg_index, Marshal.SizeOf(value), h.AddrOfPinnedObject());
            h.Free();
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
        }

        public SizeT GetWorkGroupSize(Device device)
        {
            SizeT ret = 0;
            ErrorCode res = clGetKernelWorkGroupInfo(Handle, device.Handle, KernelWorkGroupInfo.Size, ret.Size, ref ret, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return ret;
        }

        public SizeT[] GetCompileWorkGroupSize(Device device)
        {
            SizeT[] ret = new SizeT[3];
            ErrorCode res = clGetKernelWorkGroupInfo(Handle, device.Handle, KernelWorkGroupInfo.CompileSize, IntPtr.Size * 3, ret, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return ret;
        }

        public ulong GetLocalMemSize(Device device)
        {
            ulong ret = 0;
            ErrorCode res = clGetKernelWorkGroupInfo(Handle, device.Handle, KernelWorkGroupInfo.LocalMemSize, sizeof(ulong), ref ret, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return ret;
        }

        public SizeT GetPreferredWorkGroupSizeMultiple(Device device)
        {
            SizeT ret = 0;
            ErrorCode res = clGetKernelWorkGroupInfo(Handle, device.Handle, KernelWorkGroupInfo.PreferredWorkGroupSizeMultiple, ret.Size, ref ret, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return ret;
        }

        public ulong GetPrivateMemSize(Device device)
        {
            ulong ret = 0;
            ErrorCode res = clGetKernelWorkGroupInfo(Handle, device.Handle, KernelWorkGroupInfo.PrivateMemSize, sizeof(ulong), ref ret, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return ret;
        }

        public static Kernel[] GetKernels(Program program)
        {
            uint num_kernels = 0;
            ErrorCode res = clCreateKernelsInProgram(program.Handle, 0, null, ref num_kernels);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            IntPtr[] kernels = new IntPtr[num_kernels];
            res = clCreateKernelsInProgram(program.Handle, num_kernels, kernels, 0);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            Kernel[] ret = new Kernel[num_kernels];
            for (long i = 0; i < num_kernels; i++)
                ret[i] = new Kernel(kernels[i]);
            return ret;
        }

        public object Clone()
        {
            ErrorCode res = clRetainKernel(Handle);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            return new Kernel(Handle);
        }

        public void Dispose()
        {
            ErrorCode res = clReleaseKernel(Handle);
            if (res != ErrorCode.Success)
                throw new OpenCLException(res);
            Handle = IntPtr.Zero;
        }
    }
}
